import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import argparse
import cv2
import os
from net import CatDogClassifier

class GradCAM:
    """Grad-CAM热力图生成器"""
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        # 注册钩子
        self.forward_handle = target_layer.register_forward_hook(self.save_activations)
        self.backward_handle = target_layer.register_backward_hook(self.save_gradients)
    
    def __del__(self):
        # 清理钩子
        self.forward_handle.remove()
        self.backward_handle.remove()
    
    def save_activations(self, module, input, output):
        self.activations = output.detach()
    
    def save_gradients(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()
    
    def generate(self, input_tensor, target_class=None):
        # 确保输入在正确设备上
        input_tensor = input_tensor.to(next(self.model.parameters()).device)
        
        # 前向传播
        output = self.model(input_tensor)
        
        # 如果没有指定目标类别，则使用预测类别
        if target_class is None:
            target_class = torch.argmax(output, dim=1).item()
        
        # 反向传播
        self.model.zero_grad()
        one_hot = torch.zeros_like(output)
        one_hot[0][target_class] = 1
        output.backward(gradient=one_hot, retain_graph=True)
        
        # 计算权重
        pooled_gradients = torch.mean(self.gradients, dim=[0, 2, 3])
        
        # 生成热力图
        for i in range(self.activations.size(1)):
            self.activations[:, i, :, :] *= pooled_gradients[i]
        
        heatmap = torch.mean(self.activations, dim=1).squeeze()
        heatmap = np.maximum(heatmap.cpu().numpy(), 0)
        
        # 归一化
        heatmap = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap) + 1e-8)
        return heatmap

def load_model(model_path, device):
    """加载训练好的模型"""
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found: {model_path}")
    
    model = CatDogClassifier().to(device)
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()  # 设置为评估模式
    return model

def preprocess_image(image_path, img_size=224, device='cpu'):
    """预处理图像，与训练时相同"""
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"Image file not found: {image_path}")
    
    try:
        image = Image.open(image_path).convert('RGB')
    except Exception as e:
        raise ValueError(f"Could not open image: {e}")
    
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                             std=[0.229, 0.224, 0.225])
    ])
    
    original_image = np.array(image)
    input_tensor = transform(image).unsqueeze(0).to(device)  # 添加batch维度并移动到设备
    
    return original_image, input_tensor

def predict_image(model, image_tensor):
    """对图像进行预测"""
    with torch.no_grad():
        output = model(image_tensor)
    
    # 获取预测结果和置信度
    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    confidence, predicted_class = torch.max(probabilities, 0)
    
    return predicted_class.item(), confidence.item(), output

def visualize_results(image, prediction, confidence, heatmap=None):
    """可视化预测结果和热力图"""
    classes = ['cat', 'dog']
    
    plt.figure(figsize=(12, 6))
    
    # 原始图像
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title(f'Original Image\nPrediction: {classes[prediction]}')
    plt.axis('off')
    
    # 添加预测结果文本
    text = f"Predicted: {classes[prediction]}\nConfidence: {confidence:.2%}"
    plt.text(10, 30, text, 
             bbox=dict(facecolor='white', alpha=0.8),
             fontsize=10)
    
    # 热力图
    if heatmap is not None:
        # 调整热力图大小与原始图像匹配
        heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))
        heatmap = np.uint8(255 * heatmap)
        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        
        # 叠加热力图
        superimposed_img = cv2.addWeighted(
            cv2.cvtColor(image, cv2.COLOR_RGB2BGR), 0.6,
            heatmap, 0.4, 0
        )
        
        plt.subplot(1, 2, 2)
        plt.imshow(cv2.cvtColor(superimposed_img, cv2.COLOR_BGR2RGB))
        plt.title('Grad-CAM Heatmap')
        plt.axis('off')
    
    plt.tight_layout()
    plt.show()

def main():
    parser = argparse.ArgumentParser(description='Cat/Dog Image Classifier with Grad-CAM')
    parser.add_argument('--image-path', type=str, default=r"data\test\cats\cat.4001.jpg", 
                       help='Path to the input image')
    parser.add_argument('--model-path', type=str, default=r'saved_models/fold4_epoch0_acc0.9525.pth', 
                       help='Path to the trained model')
    parser.add_argument('--no-cam', action='store_true', 
                       help='Disable Grad-CAM visualization')
    args = parser.parse_args()

    try:
        # 设备设置
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"Using device: {device}")

        # 加载模型
        model = load_model(args.model_path, device)
        print(f"Loaded model from: {args.model_path}")

        # 预处理图像（确保输入张量在正确设备上）
        original_image, input_tensor = preprocess_image(args.image_path, device=device)
        
        # 进行预测
        predicted_class, confidence, output = predict_image(model, input_tensor)
        
        # 输出结果
        classes = ['cat', 'dog']
        print("\nPrediction Result:")
        print(f"Image: {args.image_path}")
        print(f"Predicted class: {classes[predicted_class]}")
        print(f"Confidence: {confidence:.2%}")

        # 生成热力图
        heatmap = None
        if not args.no_cam:
            try:
                # 获取目标层（ResNet的最后一个卷积层）
                target_layer = model.backbone.layer4[-1].conv2
                
                # 创建Grad-CAM
                grad_cam = GradCAM(model, target_layer)
                heatmap = grad_cam.generate(input_tensor, predicted_class)
            except Exception as e:
                print(f"Could not generate heatmap: {e}")
                heatmap = None

        # 可视化结果
        visualize_results(original_image, predicted_class, confidence, heatmap)
        
    except Exception as e:
        print(f"Error: {e}")
        return 1
    
    return 0

if __name__ == '__main__':
    main()